﻿// RiProgram_DISM.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <iostream>
#include <string>
#include <sstream>

#include "RrConfig.h"

#include "Q:\includes\FindPartition_Common.h"
#pragma comment(lib,"Q:/build/lib/lib/FindPartition_Common.lib")

#include "Q:\includes\ThrowError_CPP.h"
#pragma comment(lib,"Q:/build/lib/lib/ThrowError_CPP.lib")

#include "wimlib.h"
#pragma comment(lib,"Q:/includes/libwim.lib")

//#include "Format.hpp"

std::wstring String2Wstring(std::string wstr);
std::string wstring2stringB(std::wstring wstr);
std::string TCHAR2STRING(TCHAR* str);
COORD GetConsoleCursorPosition(HANDLE hConsoleOutput);
wimlib_progress_status WimlibCallBackStatus(
    wimlib_progress_msg msg_type,
    wimlib_progress_info* info,
    void* progctx
);
std::pair<int, std::string> run_cmd(const char* cmd);
std::pair<bool, std::string> FormatPartitionNTFS(char Letter);
std::string Char2Str(char c);
std::string LongDouble_2_2f(long double num);

int main()
{
    //设置标题
    SetConsoleTitleA("AZPEP SysRi Tool - Reinstall Program");

    //启动log服务
    bool Status = LogProtocol::StartLogService("AZPEP SysRi Tool - ReinstallProgram");

    if (!Status) {
        //失败，报错退出，不写入日志
        ShowError(
            __RIPROG_START_LOG_SERVICES_FAIL,
            "Fail to start log service.\nGetLastError()=" + std::to_string(GetLastError()),
            false
        );

        EmergencyExit(1);
    }
    std::string LogPath = LogProtocol::GetLogFileFullPath();

    LogProtocol::WriteLog::Info("Log start recording!");
    LogProtocol::WriteLog::Info("log file path:" + LogPath);
    LogProtocol::WriteLog::Info("");

    // 开始初始化
    LogProtocol::WriteLog::Info("Starting to obtain partitions");

    // 获取各分区
    ErrorCode FindPartitionStatus = SUCCESS;
    //std::tuple<MainProgramPartitionLetter, ReinstallPartitionLetter, EfiPartitionLetter> Partition = RiProgram_Common_GetPartition(FindPartitionStatus);

    MainProgramPartitionLetter mainProgramPartitionLetter = GetMainProgramLetter(FindPartitionStatus);
    if (FindPartitionStatus != SUCCESS) {
        ShowError(
            __RIPROG_ERROR,
            "Get mainProgramPartitionLetter Error"
        );
        EmergencyExit(1);
    }
    LogProtocol::WriteLog::Info("Image File Partition Letter:"+ Char2Str(mainProgramPartitionLetter));

    ReinstallPartitionLetter reinstallPartitionLetter = GetReinstallPartitionLetter(FindPartitionStatus);
    if (FindPartitionStatus != SUCCESS) {
        ShowError(
            __RIPROG_ERROR,
            "Get reinstallPartitionLetter Error"
        );
        EmergencyExit(1);
    }
    LogProtocol::WriteLog::Info("Reinstall System Partition Letter:" + Char2Str(reinstallPartitionLetter));

    EfiPartitionLetter efiPartitionLetter = GetEfiPartitionLetter(FindPartitionStatus);
    if (FindPartitionStatus != SUCCESS) {
        ShowError(
            __RIPROG_ERROR,
            "Get efiPartitionLetter Error"
        );
        EmergencyExit(1);
    }
    LogProtocol::WriteLog::Info("EFI Partition Letter:" + Char2Str(efiPartitionLetter));

    if (mainProgramPartitionLetter == '\0') {
        ShowError(
            __RIPROG_NO_MAINPROG_PARTITION_LETTER,
            ""
        );

        EmergencyExit(1);
    }
    else if (reinstallPartitionLetter == '\0') {
        ShowError(
            __RIPROG_NO_SYSTEM_PARTITION_LETTER,
            ""
        );

        EmergencyExit(1);
    }
    else if (efiPartitionLetter == '\0') {
        ShowError(
            __RIPROG_NO_EFI_PARTITION_LETTER,
            ""
        );

        EmergencyExit(1);
    }

    LogProtocol::WriteLog::Info("Start Read \"config.ini\"");

    //读取配置文件
    std::string InstallFilePath = "";       //映像文件相对路径
    std::string InstallFileFullPath = "";   //映像文件绝对路径
    int Index = -1;                         //索引号

    bool DisableWindowsDefender = false;    //禁用“Windows Defender”
    bool DisableWindowsUpdate = false;      //禁用Windows自动更新
    bool SkipMicrosoftLogin = false;        //跳过Microsoft登陆验证


    try {
        rr::RrConfig config;
        bool ret = config.ReadConfig(Char2Str(mainProgramPartitionLetter) +":\\config.ini");

        //读取失败，抛出异常
        if (!ret) throw "Read \"config.ini\" fail!";

        //读取ini文件
        InstallFilePath = config.ReadString("config", "InstallFile", "");
        InstallFileFullPath = Char2Str(mainProgramPartitionLetter) + ":\\" + InstallFilePath;
        Index = config.ReadInt("config", "ImageIndex", 0);

        if (config.ReadInt("config", "DisableWindowsDefender", 0) == 1) DisableWindowsDefender = true;
        if (config.ReadInt("config", "DisableWindowsUpdate", 0) == 1) DisableWindowsUpdate = true;
        if (config.ReadInt("config", "SkipMicrosoftLogin", 0) == 1) SkipMicrosoftLogin = true;
    }
    catch (std::string err) {
        //读取失败，报错退出
        ShowError(
            __RIPROG_READ_CONFIG_FAIL,
            err
        );

        EmergencyExit(1);
    }

    LogProtocol::WriteLog::Info("Opening Image File");
    //打开映像文件
    WIMStruct* WimFileStruct;       //映像文件结构体

    wimlib_error_code LoadImageStatus = 
        (wimlib_error_code)wimlib_open_wim_with_progress(
            String2Wstring(InstallFileFullPath).c_str(),        //映像路径
            NULL,                                               //以只读方式打开
            &WimFileStruct,                                     //指向WimFileStruct结构体
            WimlibCallBackStatus,
            NULL
    );

    //打开失败，报错退出
    if (LoadImageStatus != 0) {
        ShowError(
            __RIPROG_OPEN_IMAGE_FILE_FAIL,
            "Open Image File\"" + InstallFileFullPath + "\"Fail!ErrorCode=" + TCHAR2STRING((TCHAR*)wimlib_get_error_string(LoadImageStatus))
        );

        EmergencyExit(1);
    }

    LogProtocol::WriteLog::Info("Formating Partition");

    //格式化磁盘
    std::pair<bool, std::string> cmdCallBack = FormatPartitionNTFS(reinstallPartitionLetter);
    if (!cmdCallBack.first) {
        //失败，报错退出
        ShowError(
            __RIPROG_FORMAT_PARTITION_FAIL,
            cmdCallBack.second
        );
        EmergencyExit(1);
    }

    //释放映像
    wimlib_error_code ApplyImageStatus = 
        (wimlib_error_code)wimlib_extract_image(
            WimFileStruct, Index,
            String2Wstring(Char2Str(reinstallPartitionLetter) + ":\\").c_str(),
            NULL
    );

    //释放失败，报错退出
    if ((int)ApplyImageStatus != 0) {
        wimlib_free(WimFileStruct);

        ShowError(
            __RIPROG_APPLY_IMAGE_FILE_FAIL,
            TCHAR2STRING((TCHAR*)wimlib_get_error_string(ApplyImageStatus))
        );
        EmergencyExit(1);
    }

    //释放wimHandle
    wimlib_free(WimFileStruct);
    
    //重新设置标题
    SetConsoleTitleA("AZPEP SysRi Tool - Reinstall Program");

    //开始创建启动项
    //std::string CmdLine = SFormat("bcdboot {0}:\Windows /s {1}: /v /f UEFI");
    std::string CmdLine = "bcdboot " + Char2Str(reinstallPartitionLetter) + ":\\Windows /S " + Char2Str(efiPartitionLetter) + ": /F UEFI /V";
    std::pair<int,std::string> cmdCallBack2 = run_cmd(CmdLine.c_str());

    LogProtocol::WriteLog::Info("\n\n" + cmdCallBack2.second);

    if (cmdCallBack2.first != 0) {
        //失败，报错退出
        ShowError(
            __RIPROG_CREATE_STARTUP_ENTRY_FAIL,
            cmdCallBack2.second
        );
        EmergencyExit(1);
    }

    //设置注册表项
    LogProtocol::WriteLog::Info("Start setting registry entries");

    
    LogProtocol::WriteLog::Info("Reinstall System Success!");

    return 0;
}

wimlib_progress_status WimlibCallBackStatus(
    wimlib_progress_msg msg_type,
    wimlib_progress_info* info,
    void* progctx
) {
    switch (msg_type) {
        //开始释放
    case WIMLIB_PROGRESS_MSG_EXTRACT_IMAGE_BEGIN:
        
        LogProtocol::WriteLog::Info("Start Extract Image");
        break;

        //释放目录
    case WIMLIB_PROGRESS_MSG_EXTRACT_FILE_STRUCTURE: {


        std::string ProgressText = "";
        long double Progress = ((long double)info->extract.current_file_count / (long double)info->extract.end_file_count) * 100;  //进度
        int PrintProgress = floor((floor(Progress) / 2));       //打印的进度

        ProgressText += "Extract Folders: ";

        ProgressText += LongDouble_2_2f(Progress) + "%";

#ifdef _DEBUG
        //LogProtocol::WriteLog::Debug("Extract Folder: {0}"+info->extract.target);
        LogProtocol::WriteLog::Debug("Extract Folder:" + std::to_string(info->extract.current_file_count) + " / " + std::to_string(info->extract.end_file_count));
        SetConsoleTitleA(ProgressText.c_str());
#else
        COORD Pos = GetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE));
        COORD NewPos;
        NewPos.X = 0;
        NewPos.Y = Pos.Y - 1;

        SetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE), NewPos);
        LogProtocol::WriteLog::Info(ProgressText);
#endif

        break;
    }
        //写入属性
    case WIMLIB_PROGRESS_MSG_EXTRACT_METADATA: {
        uint64_t DoneFileCount = info->extract.current_file_count;
        uint64_t TotalFileCount = info->extract.end_file_count;
        //LogProtocol::WriteLog::Debug(SFormat("Write MetaData {0} / {1}  {2}%", DoneFileCount, TotalFileCount, (float)(DoneFileCount / TotalFileCount) * 100));
        LogProtocol::WriteLog::Debug("Write MetaData " + std::to_string(DoneFileCount) + " / " + std::to_string(TotalFileCount) + " " + std::to_string((long double)((long double)DoneFileCount / (long double)TotalFileCount) * 100) + " % ");
        break;

        std::string ProgressText = "";
        long double Progress = ((long double)DoneFileCount / (long double)TotalFileCount) * 100;  //进度
        int PrintProgress = floor((floor(Progress) / 2));       //打印的进度

        ProgressText += "Write MetaData: ";

        ProgressText += LongDouble_2_2f(Progress) + "%       ";

#ifdef _DEBUG
        //LogProtocol::WriteLog::Debug(SFormat("Write MetaData {0} / {1}  {2}%", DoneFileCount, TotalFileCount, (float)(DoneFileCount / TotalFileCount) * 100));
        LogProtocol::WriteLog::Debug("Write MetaData " + std::to_string(DoneFileCount) + " / " + std::to_string(TotalFileCount) + " " + std::to_string((long double)((long double)DoneFileCount / (long double)TotalFileCount) * 100) + " % ");
        SetConsoleTitleA(ProgressText.c_str());
#else
        COORD Pos = GetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE));
        COORD NewPos;
        NewPos.X = 0;
        NewPos.Y = Pos.Y - 1;

        SetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE), NewPos);
        LogProtocol::WriteLog::Info(ProgressText);
#endif
    }
    case WIMLIB_PROGRESS_MSG_WRITE_METADATA_BEGIN:
        SetConsoleTitleA("\0");
        LogProtocol::WriteLog::Info("Start Writing MetaData");
        break;

    case WIMLIB_PROGRESS_MSG_WRITE_METADATA_END:
        LogProtocol::WriteLog::Info("The release directory is complete");
        break;

        //释放文件
    case WIMLIB_PROGRESS_MSG_EXTRACT_STREAMS:{
        //std::wstring ExtractFileW(info->extract.target);
        //auto FileName = wstring2stringB(ExtractFileW);       //释放文件名

        uint64_t done_bytes = info->extract.completed_bytes;    //已释放的文件大小
        uint64_t total_bytes = info->extract.total_bytes;       //总大小

        uint64_t done_file_count = info->extract.completed_streams;    //已处理的文件数
        uint64_t total_file_count = info->extract.total_streams;       //文件总数

        //LogProtocol::WriteLog::Debug(SFormat("Extract File:{0}", FileName));
        // 
        //LogProtocol::WriteLog::Debug(SFormat("Files processed: {0} / {1}", done_file_count, total_file_count));
        //LogProtocol::WriteLog::Debug(SFormat("File data has been extracted: {0} Byte(s) / {1} Byte(s)", done_bytes, total_bytes));
        LogProtocol::WriteLog::Debug("Files processed:" + std::to_string(done_file_count) + " / " + std::to_string(total_file_count));
        LogProtocol::WriteLog::Debug("File data has been extracted: " + std::to_string(done_bytes) + " Byte(s) / " + std::to_string(total_bytes) + " Byte(s)");
        LogProtocol::WriteLog::Debug("\n");

        //break;

        // 输出进度

        std::string ProgressText = "";

        long double Progress = ((long double)done_bytes / (long double)total_bytes) * 100;  //进度

        int PrintProgress = (int)floor((int)floor(Progress) / 4);       //打印的进度

        ProgressText += "Extract Files: ";

        ProgressText += LongDouble_2_2f(Progress) + "%       ";

#ifdef _DEBUG
        SetConsoleTitleA(ProgressText.c_str());
#else
        COORD Pos = GetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE));
        COORD NewPos;
        NewPos.X = 0;
        NewPos.Y = Pos.Y - 1;

        SetConsoleCursorPosition(GetStdHandle(STD_OUTPUT_HANDLE), NewPos);
        LogProtocol::WriteLog::Info(ProgressText);
#endif

        break;
        }
    case WIMLIB_PROGRESS_MSG_EXTRACT_IMAGE_END:
        LogProtocol::WriteLog::Info("Extract Image Done!\n");
        break;
    }
    return WIMLIB_PROGRESS_STATUS_CONTINUE;

}

std::wstring String2Wstring(std::string wstr)
{
    std::wstring res;
    int len = MultiByteToWideChar(CP_ACP, 0, wstr.c_str(), wstr.size(), nullptr, 0);
    if (len < 0) {
        return res;
    }
    wchar_t* buffer = new wchar_t[len + 1];
    if (buffer == nullptr) {
        return res;
    }
    MultiByteToWideChar(CP_ACP, 0, wstr.c_str(), wstr.size(), buffer, len);
    buffer[len] = '\0';
    res.append(buffer);
    delete[] buffer;
    return res;
}


//将wstring转换成string  
std::string wstring2stringB(std::wstring wstr)
{
    std::string result;
    //获取缓冲区大小，并申请空间，缓冲区大小事按字节计算的  
    int len = WideCharToMultiByte(CP_ACP, 0, wstr.c_str(), wstr.size(), NULL, 0, NULL, NULL);
    char* buffer = new char[len + 1];
    //宽字节编码转换成多字节编码  
    WideCharToMultiByte(CP_ACP, 0, wstr.c_str(), wstr.size(), buffer, len, NULL, NULL);
    buffer[len] = '\0';
    //删除缓冲区并返回值  
    result.append(buffer);
    delete[] buffer;
    return result;

}

std::string TCHAR2STRING(TCHAR* str)
{
    std::string strstr;
    try
    {
        int iLen = WideCharToMultiByte(CP_ACP, 0, str, -1, NULL, 0, NULL, NULL);

        char* chRtn = new char[iLen * sizeof(char)];

        WideCharToMultiByte(CP_ACP, 0, str, -1, chRtn, iLen, NULL, NULL);

        strstr = chRtn;
    }
    catch (std::exception e)
    {
    }

    return strstr;
}

COORD GetConsoleCursorPosition(HANDLE hConsoleOutput)
{
    CONSOLE_SCREEN_BUFFER_INFO cbsi;
    if (GetConsoleScreenBufferInfo(hConsoleOutput, &cbsi))
    {
        return cbsi.dwCursorPosition;
    }
    else
    {
        // The function failed. Call GetLastError() for details.
        COORD invalid = { 0, 0 };
        return invalid;
    }
}

/// <summary>
/// 快速格式化NTFS分区
/// </summary>
/// <param name="Letter">盘符，如“C”</param>
/// <returns>成功返回true，失败返回false</returns>
std::pair<bool, std::string> FormatPartitionNTFS(char Letter) {
    //生成命令
    std::string CommandLine = "format " + Char2Str((Letter)) + ": /V:Windows /Q /X /Y";

    int ReturnCode = 0;//退出代码
    std::string ReturnText = "";		//返回信息

    //调用cmd
    std::pair<int, std::string> cmdCall = run_cmd(CommandLine.c_str());

    //返回
    if (cmdCall.first != 0) return { false,CommandLine }; else return { true,CommandLine };
}

std::pair<int, std::string> run_cmd(const char* cmd)
{
    char MsgBuff[1024];
    int MsgLen = 1020;
    FILE* fp;
    std::string ReturnText = "";

    if (cmd == NULL)
    {
        return { -1,ReturnText };
    }
    if ((fp = _popen(cmd, "r")) == NULL)
    {
        return { -2,ReturnText };
    }
    else
    {
        memset(MsgBuff, 0, MsgLen);

        //读取命令执行过程中的输出
        while (fgets(MsgBuff, MsgLen, fp) != NULL)
        {
            std::string a(MsgBuff);
            ReturnText += a + "\n";
        }

        //关闭执行的进程
        int returnCode = _pclose(fp);

        return { returnCode,ReturnText };
    }
}

std::string Char2Str(char c) {
    std::string s(1, c);
    return s;
}

std::string LongDouble_2_2f(long double num) {
    char c[40];

    sprintf_s(c, 40 ,"%.2f", num);

    std::stringstream ss;
    ss << c;

    return ss.str();
}